import time
import yaml
import nacos
import shlex
import atexit
import yagmail
import schedule
import subprocess
import multiprocessing

from typing import Dict
from pathlib import Path
from loguru import logger
from datetime import datetime
from threading import Lock, Timer
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, FileResponse
from sqlmodel import Field, Session, SQLModel, create_engine, select
from fastapi import FastAPI, File, Form, UploadFile, BackgroundTasks, status

# 创建 FastAPI 实例
app = FastAPI()
# 设置跨域资源共享
app.add_middleware(
    CORSMiddleware,
    allow_origins=[
        'http://localhost',
        'http://localhost:5001',
        'http://10.181.1.92',
        'http://10.181.1.92:5001'
    ],
    allow_credentials=True,
    allow_methods=['*'],
    allow_headers=['*'],
)

# 任务集合
task_set = set()
# 最大任务数量
max_size = 2
# 任务锁
task_lock = Lock()

# 文件目录
file_dir = Path(__file__).parent / Path('files')
if not file_dir.is_dir():
    file_dir.mkdir(parents=True, exist_ok=True)

# 日志目录
log_dir = Path(__file__).parent / Path('logs')
if not log_dir.is_dir():
    log_dir.mkdir(parents=True, exist_ok=True)

# 添加日志
log_path = log_dir / Path('app.log')
logger.add(
    sink=log_path,
    level='INFO',
    rotation="100 MB",
    retention='10 days'
)

# Nacos客户端
client = nacos.NacosClient(
    server_addresses='http://10.184.17.100:8848',
    namespace='public'
)

# 创建数据库引擎
db_url = 'mysql+pymysql://root:123456@10.184.17.100:3306/ry-material'
# db_url = 'mysql+pymysql://root:root@127.0.0.1:3306/db_hotel'
engine = create_engine(db_url)


class Task(SQLModel, table=True):
    """
    任务表
    """
    id: int = Field(default=None, primary_key=True, nullable=False, description='自增ID')
    name_magnetic: str | None = Field(default=None, max_length=255, nullable=True, description='电磁数据名称')
    create_time_magnetic: datetime | None = Field(default=None, nullable=True, description='电磁仿真创建时间')
    flux_density_mod_image: str | None = Field(default=None, nullable=True, description='磁通密度模图片 (base64)')
    lorentz_force_r_image: str | None = Field(default=None, nullable=True, description='洛伦兹力r方向图片 (base64)')
    lorentz_force_z_image: str | None = Field(default=None, nullable=True, description='洛伦兹力z方向图片 (base64)')
    copper_loss_image: str | None = Field(default=None, nullable=True, description='铜耗图片 (base64)')
    flux_density_mod_txt: str | None = Field(default=None, max_length=255, nullable=True, description='磁通密度模txt文件地址')
    lorentz_force_r_txt: str | None = Field(default=None, max_length=255, nullable=True, description='洛伦兹力r方向txt文件地址')
    lorentz_force_z_txt: str | None = Field(default=None, max_length=255, nullable=True, description='洛伦兹力z方向txt文件地址')
    copper_loss_txt: str | None = Field(default=None, max_length=255, nullable=True, description='铜耗txt文件地址')
    type_heat: int | None = Field(default=None, nullable=True, description='热学仿真类型-正常/过载/堵塞')
    inputfile_loss_heat: str | None = Field(default=None, max_length=255, nullable=True, description='热学输入损耗文件地址')
    x1: float | None = Field(default=None, nullable=True, description='堵塞位置矩形对角线x1坐标')
    y1: float | None = Field(default=None, nullable=True, description='堵塞位置矩形对角线y1坐标')
    x2: float | None = Field(default=None, nullable=True, description='堵塞位置矩形对角线x2坐标')
    y2: float | None = Field(default=None, nullable=True, description='堵塞位置矩形对角线y2坐标')
    temperature_heat: str | None = Field(default=None, max_length=255, nullable=True, description='输出温度csv文件地址')
    image_blockage_heat: str | None = Field(default=None, nullable=True, description='堵塞位置图片的base64编码内容')
    image1_temperature_heat: str | None = Field(default=None, nullable=True, description='图片1温度的base64编码内容')
    image2_vector_heat: str | None = Field(default=None, nullable=True, description='图片2矢量的base64编码内容')
    magnetic_force_r_txt: str | None = Field(default=None, max_length=255, nullable=True, description='R径向磁力分量txt文件地址')
    magnetic_force_z_txt: str | None = Field(default=None, max_length=255, nullable=True, description='Z轴向磁力分量txt文件地址')
    mechanics_result_frd: str | None = Field(default=None, max_length=255, nullable=True, description='力学仿真结果frd文件地址')
    mechanics_result_vtk: str | None = Field(default=None, max_length=255, nullable=True, description='力学仿真结果vtk文件地址')
    mechanics_result_vtp: str | None = Field(default=None, max_length=255, nullable=True, description='力学仿真结果vtp文件地址')
    magnetic_status: int | None = Field(default=0, nullable=True, description='电磁仿真任务状态（0-未开始，1-正在执行，2-执行成功，3-执行失败）')
    heat_status: int | None = Field(default=0, nullable=True, description='热学仿真任务状态（0-未开始，1-正在执行，2-执行成功，3-执行失败）')
    mechanics_status: int | None = Field(default=0, nullable=True, description='力学仿真任务状态（0-未开始，1-正在执行，2-执行成功，3-执行失败）')


def run_command(command: str, cwd: Path = None) -> bool:
    """
    执行命令
    """
    result = False
    try:
        completed_process = subprocess.run(
            args=shlex.split(command),
            capture_output=True,
            text=True,
            cwd=cwd
        )
        logger.info('命令执行成功：' + str(completed_process))
        if completed_process.returncode == 0:
            result = True
    except Exception as e:
        logger.error('命令执行失败：' + str(repr(e)))

    return result


@app.get('/')
@app.get('/api')
@app.get('/api/health')
def health() -> str:
    """
    服务健康状态
    """
    return 'alive'


@app.post('/api/download', response_model=None)
async def download(
        file_path: str = Form(default=None)
) -> JSONResponse | FileResponse:
    """
    下载服务器文件
    """
    if file_path is None:
        return JSONResponse(
            status_code=status.HTTP_400_BAD_REQUEST,
            content={
                'code': status.HTTP_400_BAD_REQUEST,
                'msg': '参数值缺失',
                'data': None
            }
        )

    # 获取请求参数
    logger.info('文件路径：' + str(file_path))

    # 检查文件是否存在
    file = Path(file_path)
    if file.exists():
        return FileResponse(
            status_code=status.HTTP_200_OK,
            path=file_path
        )
    else:
        return JSONResponse(
            status_code=status.HTTP_404_NOT_FOUND,
            content={
                'code': status.HTTP_404_NOT_FOUND,
                'msg': '文件不存在',
                'data': None
            }
        )


@app.post('/api/execute_v3', response_model=None)
async def execute_v3(
        magnetic_force_r_txt_file: UploadFile = File(default=None),
        magnetic_force_z_txt_file: UploadFile = File(default=None),
        task_id: str = Form(default=None),
        email_address: str = Form(default=None),
        background_tasks: BackgroundTasks = None
) -> JSONResponse:
    """
    执行程序和脚本进行力学仿真
    """
    if magnetic_force_r_txt_file is None or magnetic_force_z_txt_file is None is None or task_id is None or email_address is None:
        return JSONResponse(
            status_code=status.HTTP_400_BAD_REQUEST,
            content={
                'code': status.HTTP_400_BAD_REQUEST,
                'msg': '参数值缺失',
                'data': None
            }
        )

    # 获取请求参数
    mechanics_result_file_name = f'110kv_struct_task_{task_id}'
    logger.info('R径向磁力分量txt文件：' + str(magnetic_force_r_txt_file.filename))
    logger.info('Z轴向磁力分量txt文件：' + str(magnetic_force_z_txt_file.filename))
    logger.info('任务ID：' + str(task_id))
    logger.info('力学仿真结果文件名称：' + str(mechanics_result_file_name))

    # 判断任务是否存在
    with Session(engine) as session:
        statement = select(Task).where(Task.id == int(task_id))
        results = session.exec(statement)
        tasks = list(results)

    if len(tasks) == 0:
        logger.error(f'任务{task_id}不存在')
        return JSONResponse(
            status_code=status.HTTP_404_NOT_FOUND,
            content={
                'code': status.HTTP_404_NOT_FOUND,
                'msg': f'任务{task_id}不存在',
                'data': None
            }
        )

    # 阻止重复提交
    with task_lock:
        if task_id in task_set:
            return JSONResponse(
                status_code=status.HTTP_403_FORBIDDEN,
                content={
                    'code': status.HTTP_403_FORBIDDEN,
                    'msg': f'任务{task_id}的力学仿真任务已经在后台执行，请不要重复提交',
                    'data': None
                }
            )

        if len(task_set) >= max_size:
            return JSONResponse(
                status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
                content={
                    'code': status.HTTP_503_SERVICE_UNAVAILABLE,
                    'msg': '力学仿真计算任务的数量超出服务器能支持的最大范围，请稍后重新提交',
                    'data': None
                }
            )

        # 将任务添加到任务集合中
        task_set.add(task_id)

    # 保存传入的输入文件
    upload_dir = file_dir / Path(str(mechanics_result_file_name))
    if not upload_dir.is_dir():
        upload_dir.mkdir(parents=True, exist_ok=True)
    magnetic_force_r_txt_file_path = upload_dir / Path(magnetic_force_r_txt_file.filename)
    magnetic_force_z_txt_file_path = upload_dir / Path(magnetic_force_z_txt_file.filename)
    with open(magnetic_force_r_txt_file_path, mode='wb+') as file_object:
        file_object.write(magnetic_force_r_txt_file.file.read())
    with open(magnetic_force_z_txt_file_path, mode='wb+') as file_object:
        file_object.write(magnetic_force_z_txt_file.file.read())

    # 添加后台任务
    background_tasks.add_task(
        simulate,
        magnetic_force_r_txt_file,
        magnetic_force_z_txt_file,
        task_id, email_address,
        mechanics_result_file_name
    )

    return JSONResponse(
        status_code=status.HTTP_202_ACCEPTED,
        content={
            'code': status.HTTP_202_ACCEPTED,
            'msg': f'任务{task_id}的力学仿真任务已经成功提交至后台执行，请检查邮箱通知',
            'data': None
        }
    )


def simulate(
        magnetic_force_r_txt_file: UploadFile = None,
        magnetic_force_z_txt_file: UploadFile = None,
        task_id: str = None,
        email_address: str = None,
        mechanics_result_file_name: str = None
) -> None:
    """
    后台执行程序和脚本进行力学仿真
    """
    logger.info(f'开始执行任务{task_id}的力学仿真任务，当前队列中的任务：{str(task_set)}')

    # 创建存放文件的目录
    upload_dir = file_dir / Path(str(mechanics_result_file_name))
    magnetic_force_r_txt_file_path = upload_dir / Path(magnetic_force_r_txt_file.filename)
    magnetic_force_z_txt_file_path = upload_dir / Path(magnetic_force_z_txt_file.filename)
    # 设置仿真结果的保存路径
    mechanics_result_frd_file_path = upload_dir / Path(f'{mechanics_result_file_name}.frd')
    mechanics_result_vtk_file_path = upload_dir / Path(f'{mechanics_result_file_name}.vtk')
    mechanics_result_vtp_file_path = upload_dir / Path(f'{mechanics_result_file_name}.vtp')

    flag = True
    email_content = f'力学仿真任务{task_id}执行成功，请及时下载'
    try:
        # 更新力学仿真输入文件和力学仿真任务状态为正在执行
        with Session(engine) as session:
            statement = select(Task).where(Task.id == int(task_id))
            results = session.exec(statement)
            task = results.one()
            task.magnetic_force_r_txt = str(magnetic_force_r_txt_file_path),
            task.magnetic_force_z_txt = str(magnetic_force_z_txt_file_path),
            task.mechanics_status = 1
            session.add(task)
            session.commit()
            session.refresh(task)

        """
        # 在原始模型基础上添加电磁力（10min）
        command1 = f'python scripts/AddMagForce.py exe/110kv_new_model.inp files/{mechanics_result_file_name}/{magnetic_force_r_txt_file.filename} files/{mechanics_result_file_name}/{magnetic_force_z_txt_file.filename} files/{mechanics_result_file_name}/{mechanics_result_file_name}.inp'
        logger.info('执行命令1：' + str(command1))
        if run_command(command1) is False:
            flag = False
            email_content = f'力学仿真任务{task_id}执行失败，{command1}命令执行出错，请重新提交'

        # 设置并行计算的线程数量（1s）
        command2 = f'setx OMP_NUM_THREADS {multiprocessing.cpu_count()}'
        logger.info('执行命令2：' + str(command2))
        if run_command(command2) is False:
            flag = False
            email_content = f'力学仿真任务{task_id}执行失败，{command2}命令执行出错，请重新提交'

        # 调用求解器完成力学仿真计算（2h40min）
        command3 = f'exe/ccx_dynamic.exe -i files/{mechanics_result_file_name}/{mechanics_result_file_name}'
        logger.info('执行命令3：' + str(command3))
        if run_command(command3) is False:
            flag = False
            email_content = f'力学仿真任务{task_id}执行失败，{command3}命令执行出错，请重新提交'

        # 将frd结果文件转换为vtk文件（4min）
        command4 = f'python scripts/frd2vtk.py files/{mechanics_result_file_name}/{mechanics_result_file_name}.frd vtk'
        logger.info('执行命令4：' + str(command4))
        if run_command(command4) is False:
            flag = False
            email_content = f'力学仿真任务{task_id}执行失败，{command4}命令执行出错，请重新提交'

        # 将vtk文件转换为vtp文件用于显示（1min）
        command5 = f'python scripts/vtk2vtp.py files/{mechanics_result_file_name}/{mechanics_result_file_name}.vtk'
        logger.info('执行命令5：' + str(command5))
        if run_command(command5) is False:
            flag = False
            email_content = f'力学仿真任务{task_id}执行失败，{command5}命令执行出错，请重新提交'
        """

        # todo: =============================================生产环境删除
        time.sleep(10)
        mechanics_result_frd_file_path = upload_dir / Path(f'{mechanics_result_file_name}.frd')
        mechanics_result_vtk_file_path = upload_dir / Path(f'{mechanics_result_file_name}.vtk')
        mechanics_result_vtp_file_path = upload_dir / Path(f'{mechanics_result_file_name}.vtp')
        with open(mechanics_result_frd_file_path, mode='wb+') as file_object:
            file_object.write('test demo frd'.encode())
        with open(mechanics_result_vtk_file_path, mode='wb+') as file_object:
            file_object.write('test demo vtk'.encode())
        with open(mechanics_result_vtp_file_path, mode='wb+') as file_object:
            file_object.write('test demo vtp'.encode())
        # todo: =============================================生产环境删除

    except Exception as e:
        logger.error(f'服务器内部错误：{str(e)}')
        flag = False
        email_content = f'力学仿真任务{task_id}执行失败，服务器内部错误，请重新提交'

    if flag:
        # 更新力学仿真结果文件和力学仿真任务状态为执行成功
        with Session(engine) as session:
            statement = select(Task).where(Task.id == int(task_id))
            results = session.exec(statement)
            task = results.one()
            task.mechanics_result_frd = str(mechanics_result_frd_file_path),
            task.mechanics_result_vtk = str(mechanics_result_vtk_file_path),
            task.mechanics_result_vtp = str(mechanics_result_vtp_file_path),
            task.mechanics_status = 2
            session.add(task)
            session.commit()
            session.refresh(task)
    else:
        # 更新力学仿真任务状态为执行失败
        with Session(engine) as session:
            statement = select(Task).where(Task.id == int(task_id))
            results = session.exec(statement)
            task = results.one()
            task.mechanics_status = 3
            session.add(task)
            session.commit()
            session.refresh(task)

    # 阻止重复提交
    with task_lock:
        # 删除任务集合中的任务
        task_set.remove(task_id)

    logger.info(f'结束执行任务{task_id}的力学仿真任务，当前队列中的任务：{str(task_set)}')

    # 发送邮件
    send_email(email_address, email_content)


def send_email(
        email_address: str, email_content: str
) -> None:
    """
    发送邮件
    """
    yag = yagmail.SMTP(
        user='1640561618@qq.com',
        password='cyncpjahlrbpeeha',
        host='smtp.qq.com',
        port=465,
        smtp_set_debuglevel=1
    )
    yag.send(
        to=email_address,
        subject='力学仿真任务提醒',
        contents=email_content,
    )
    yag.send_unsent()
    yag.close()


def service_register() -> bool:
    """
    注册服务
    """
    response = client.add_naming_instance(
        service_name='ruoyi-flask',
        ip='0.0.0.0',
        port=5001
    )
    if response:
        logger.info('向Nacos治理中心注册微服务成功：' + str(response))
    else:
        logger.error('向Nacos治理中心注册微服务失败：' + str(response))

    return response


def service_configuration() -> Dict:
    """
    获取服务配置
    """
    configuration = yaml.full_load(
        stream=client.get_config(
            data_id='ruoyi-flask-dev.yml',
            group='DEFAULT_GROUP'
        )
    )
    logger.info('从Nacos治理中心获取微服务的配置：' + str(configuration))

    return configuration


@schedule.repeat(schedule.every(60).seconds)
def service_heartbeat() -> None:
    """
    每隔60s发送一次心跳
    """
    try:
        response = client.send_heartbeat(
            service_name='ruoyi-flask',
            ip='0.0.0.0',
            port=5001
        )
    except Exception as e:
        logger.error('向Nacos治理中心发送心跳请求失败：' + str(repr(e)))


def start_scheduler_thread() -> None:
    """
    启动定时器
    """
    while True:
        schedule.run_pending()
        time.sleep(1)


def clear_on_exit() -> None:
    """
    退出时清除所有任务
    """
    task_set.clear()


if __name__ == '__main__':
    atexit.register(clear_on_exit)
    Timer(interval=5, function=service_register).start()
    Timer(interval=10, function=service_configuration).start()
    Timer(interval=15, function=start_scheduler_thread).start()
